import numpy as np
import matplotlib.pyplot as plt
import torch

from eval import show_without_plt

# 中文显示
plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False

def show_plot(data_1, data_2):

    # x 轴的刻度位置
    x = np.arange(len(data_1))  # [0, 1, 2, 3, 4]

    # 柱子宽度
    width = 0.4

    # 创建图形
    plt.figure(figsize=(20, 6))

    # 绘制两组柱子
    bars1 = plt.bar(x - width / 2, data_1, width=width, label="旋转前", color='skyblue', edgecolor='black')
    bars2 = plt.bar(x + width / 2, data_2, width=width, label="旋转后", color='salmon', edgecolor='black')

    # 添加标题和坐标轴标签
    plt.title("每个block的feature map稀疏度", fontsize=16)
    plt.xlabel("block", fontsize=12)
    plt.ylabel("稀疏度", fontsize=12)

    # 设置 x 轴刻度和标签
    categories = [f"block{i+1}" for i in range(len(data_1))]  # 每个柱子的类别标签


    plt.xticks(x, categories, fontsize=10)

    # 添加图例
    plt.legend(fontsize=12)

    # 在柱子顶部添加数值标签
    for bars in [bars1, bars2]:
        for bar in bars:
            height = bar.get_height()
            plt.text(bar.get_x() + bar.get_width() / 2, height,
                     f'{height*100:.1f}%', ha='center', va='bottom', fontsize=10)

    # 添加网格线
    plt.grid(axis='y', linestyle='--', linewidth=0.5)

    # 调整布局，防止元素重叠
    plt.tight_layout()

    save_path = f"每个block的feature map稀疏度"
    if save_path:
        plt.savefig(save_path, dpi=300)  # 保存为高分辨率图片
        print(f"Figure saved to {save_path}")

    # 显示图形
    plt.show()

def show_plot_single(data_1, data_2):
    # x 轴的刻度位置
    x = np.array([0])
    # 柱子宽度
    width = 0.3
    # 创建图形
    plt.figure(figsize=(6,6))

    # 绘制两组柱子
    bars1 = plt.bar(x - width / 2, data_1, width=width, label="旋转前", color='skyblue', edgecolor='black')
    bars2 = plt.bar(x + width / 2, data_2, width=width, label="旋转后", color='salmon', edgecolor='black')

    # 添加标题和坐标轴标签
    plt.title("全局平均稀疏度", fontsize=16)
    plt.xlabel("block", fontsize=12)
    plt.ylabel("稀疏度", fontsize=12)

    # 设置 x 轴刻度和标签
    categories = ['']# 每个柱子的类别标签

    plt.xticks(x, categories, fontsize=10)

    # 添加图例
    plt.legend(fontsize=12)

    # 在柱子顶部添加数值标签
    for bars in [bars1, bars2]:
        for bar in bars:
            height = bar.get_height()
            plt.text(bar.get_x() + bar.get_width() / 2, height,
                     f'{height * 100:.2f}%', ha='center', va='bottom', fontsize=18)

    # 添加网格线
    plt.grid(axis='y', linestyle='--', linewidth=0.5)

    # 调整布局，防止元素重叠
    plt.tight_layout()

    save_path = f"全局平均稀疏度"
    if save_path:
        plt.savefig(save_path, dpi=300)  # 保存为高分辨率图片
        print(f"Figure saved to {save_path}")

    # 显示图形
    plt.show()

if __name__ == '__main__':
    data_1 = []
    data_2 = []

    load_tensors_1 = torch.stack(
        [torch.load(f"../ffn/second_line_input/tensor{i}.pth", weights_only=True) for i in range(32)])
    load_tensors_2 = torch.matmul(load_tensors_1, torch.load("../1-1-2_target_657.pth", weights_only=True).to("cpu"))

    for i in range(32):
        temp1 = show_without_plt(load_tensors_1[i])
        temp2 = show_without_plt(load_tensors_2[i])
        data_1.append(temp1["lower_ratio_1e_4"])
        data_2.append(temp2["lower_ratio_1e_4"])

    show_plot(data_1, data_2)

    data_1_avg = sum(data_1) / len(data_1)
    data_2_avg = sum(data_2) / len(data_2)

    print(f"旋转前平均稀疏度：{data_1_avg}")
    print(f"旋转后平均稀疏度：{data_2_avg}")

    show_plot_single([data_1_avg],[data_2_avg])






